package cl.regalloc;

import java.io.BufferedReader;
import java.io.FileReader;
import java.util.HashMap;
import java.util.HashSet;
import java.util.LinkedList;
import java.util.List;

import cl.regalloc.graph.InterferenceGraph;
import cl.regalloc.graph.Node;
import cl.regalloc.parse.BasicBlock;
import cl.regalloc.parse.BasicBlockFactory;
import cl.regalloc.parse.Instruction;
import cl.regalloc.parse.Label;
import cl.regalloc.parse.LabelFactory;
import cl.regalloc.parse.Register;
import cl.regalloc.parse.RegisterFactory;
import cl.regalloc.parse.Routine;

public class RegisterAllocation {
	private static final Integer AVAILABLE_REGISTERS = 32;

	private static LinkedList<String> params = new LinkedList<String>();
	private static LinkedList<Routine> program = new LinkedList<Routine>();


	private static void printUsage() {
		System.out.println("java RegisterAllocation arg0 arg1");
		System.out.println("arg0 - input assembly code file name");
		System.out.println("arg1 - output assembly code file name");
	}

	private static InterferenceGraph buildGraph(String inputFile) {
		try {
			BufferedReader fileReader = new BufferedReader(new FileReader(inputFile));
			String currentLine = fileReader.readLine();
			if (currentLine != null) currentLine = currentLine.trim();
			//			boolean moveInstruction = false;
			while(currentLine != null) {
				if (currentLine.contains("#") && (!currentLine.contains(".word"))) {
					/* ignore all the comments, EXCEPT those with .word representing machine coded instruction */
					System.out.printf("comment detectect: %s%n", currentLine);
					currentLine = fileReader.readLine();
					if (currentLine != null) currentLine = currentLine.trim();
					continue;
				}

				if (currentLine.contains(".set") || currentLine.contains(".text") || currentLine.equals("")) {
					params.add(currentLine);
					System.out.printf(".set or .text or \"\" detectect, storing to parameters: %s%n", currentLine);
					currentLine = fileReader.readLine();
					if (currentLine != null) currentLine = currentLine.trim();
					continue;
				}

				if (currentLine.contains(".ent") && currentLine.contains("codeseg")) {
					System.out.printf("standard __codeseg detectect: %s %n storing to params%n%n", currentLine);
					/* default __codeseg, just store it */
					params.add(currentLine);
					currentLine = fileReader.readLine();
					if (currentLine != null) currentLine = currentLine.trim();
					while (!currentLine.contains(".end")) {
						params.add(currentLine);
						currentLine = fileReader.readLine();
						if (currentLine != null) currentLine = currentLine.trim();
					}
					params.add(currentLine);	/* .end __codeseg */
					currentLine = fileReader.readLine();
					if (currentLine != null) currentLine = currentLine.trim();
					continue;
				}

				if (currentLine.contains(".ent")) {
					System.out.printf("new routine detectect: %s %n storing to params%n%n", currentLine);
					params.add(currentLine);	/* save entry code */
					currentLine = currentLine.replace(".ent", "").trim();
					String currentRoutineName = currentLine.substring(0, currentLine.indexOf(","));	/* extract name of routine */

					System.out.printf("routine name \"%s\" exracted%n", currentRoutineName);
					LinkedList<BasicBlock> currentBasicBlockList = new LinkedList<BasicBlock>();
					BasicBlock currentBasicBlock = new BasicBlock();

					currentLine = fileReader.readLine();
					if (currentLine != null) currentLine = currentLine.trim();
					while (!currentLine.contains(".end")) {
						if (currentLine == null) break;
						if (currentLine.contains(".globl")) {
							/* handle .globl case */
							System.out.printf(".gbol detectect%n", currentLine);
							currentLine = fileReader.readLine();
							if (currentLine != null) currentLine = currentLine.trim();
							continue;
						}

						if (currentLine.contains(".data")) {
							/* handle .data case */
							System.out.printf(".data detectect: %s %n storing to params%n%n", currentLine);
							currentLine = fileReader.readLine();
							if (currentLine != null) currentLine = currentLine.trim();
							continue;
						}

						if (currentLine.contains(":")) {
							/* handle labels */

							String currentLabelName = currentLine.substring(0, currentLine.indexOf(":"));
							System.out.printf("label \"%s\" detectect: %s %n breaking basic block%n%n", currentLabelName, currentLine);
							Label currentLabel = LabelFactory.getLabel(currentLabelName);
							if (!currentBasicBlock.isEmpty()) {
								/* new label always starts a new basic block */
								currentBasicBlockList.add(currentBasicBlock);

								BasicBlock newBasicBlock = BasicBlockFactory.getBasicBlockWithLabel(currentLabel);
								currentBasicBlock.addSucc(newBasicBlock);
								newBasicBlock.addPrdc(currentBasicBlock);
								currentBasicBlock = newBasicBlock;


							} else {
								BasicBlock labeledBasicBlock = BasicBlockFactory.getBasicBlockWithLabel(currentLabel);		
								// transfer predecessors and successors, cannot create new BasicBlock with label directly
								// because it may be created already by factory!
								currentBasicBlock = BasicBlockFactory.transferLinks(currentBasicBlock, labeledBasicBlock);
							}
							currentLine = fileReader.readLine();
							if (currentLine != null) currentLine = currentLine.trim();
							continue;
						} 

						if (currentLine.contains(".word")) {
							System.out.printf(".word detectect: %s%n extracting instruction%n", currentLine);
							/* instruction in machine code, read instruction following "#" */
							if (currentLine.contains("#")) {
								System.out.printf("pos %d%n",	currentLine.indexOf("#"));
								currentLine = currentLine.substring(currentLine.indexOf("#") + 2, currentLine.length());
								currentLine = currentLine.trim();
								System.out.printf(".word detectect: %s extracted%n%n", currentLine);
							} else {
								throw new Exception();
							}
						} /* fall through to process as normal instruction */

						/* normal instruction */
						System.out.printf("instruction \"%s\" detected%n", currentLine);
						Instruction currentInstruction = new Instruction(currentLine);
						currentLine = fileReader.readLine();
						if (currentLine != null) currentLine = currentLine.trim();

						if (currentInstruction.isConditional() || currentInstruction.isJump()) {
							/* needs to handle delay slot */
							System.out.printf("Branch or Jump detected, handling delay slot...%n");
							Instruction delaySlotInstruction = new Instruction(currentLine);
							currentBasicBlock.insertInstruction(delaySlotInstruction);
							currentBasicBlock.insertInstruction(currentInstruction);	

							/* new conditional or jump always start a new basic block */
							currentBasicBlockList.add(currentBasicBlock);
							if (currentInstruction.hasLabelDst()) {
								Label dstLabel = currentInstruction.getDst();
								BasicBlock labeledBasicBlock = BasicBlockFactory.getBasicBlockWithLabel(dstLabel);	/* for the purpose of linking prdcs and succs */
								currentBasicBlock.addSucc(labeledBasicBlock);
								labeledBasicBlock.addPrdc(currentBasicBlock);
							}

							BasicBlock newBasicBlock = new BasicBlock();
							currentBasicBlock.addSucc(newBasicBlock);
							newBasicBlock.addPrdc(currentBasicBlock);
							currentBasicBlock = newBasicBlock;

							currentLine = fileReader.readLine();
							if (currentLine != null) currentLine = currentLine.trim();

						} else {
							currentBasicBlock.insertInstruction(currentInstruction);
						}
					} /* while (!currentLine.contains(".end")) */

					/* now store the routine */
					if (!currentLine.contains(currentRoutineName)) throw new Exception();	/* making sure it's the same method */
					params.add(currentLine);
					Routine newRoutine = new Routine(currentRoutineName, currentBasicBlockList);
					program.add(newRoutine); /* add to the program list */

					currentLine = fileReader.readLine();
					if (currentLine != null) currentLine = currentLine.trim();
					continue;
				} /* if (currentLine.contains(".ent")) */
			} /* while(currentLine != null) */
		} catch (Exception e) {
			e.printStackTrace();
		}


		/* BASIC BLOCK PRINTING */
		System.out.printf("%n%n** BASIC BLOCK PRINTOUT **%n");
		for (Routine currentRoutine : program)
			for (BasicBlock currentBasicBlock : currentRoutine.getRoutine()) {
				System.out.printf("%ncurrentBasicBlock: %s%n", currentBasicBlock.toString());
				if (currentBasicBlock.getLabel() != null)
					System.out.printf("Basic Block LABEL: %s%n", currentBasicBlock.getLabel().toString());
				for (Instruction currentInstruction : currentBasicBlock.getBlock())
					System.out.printf("Instruction: %s%n", currentInstruction.toString());
				for (BasicBlock currentPrdc : currentBasicBlock.getPrdcs())
					System.out.printf("Prdc: %s%n", currentPrdc);
				for (BasicBlock currentSucc : currentBasicBlock.getSuccs())
					System.out.printf("Succ: %s%n", currentSucc);
			}
		System.out.printf("%n%n** END OF BASIC BLOCK PRINTOUT **%n");

		System.out.printf("Number of registers used in input: %d%n", RegisterFactory.getRegisterKeys().size());
		convertToSsa();	// SSA conversion
		performLva();	// Live Variable Analysis
		return generateInterferenceGraph();

	}


	private static boolean convertToSsa() {
		HashMap<Register, Register> regMap = new HashMap<Register, Register>();

		System.out.printf("%n%nSSA conversion begins...%n");

		for (Routine currentRoutine : program) {
			System.out.printf("Routine: %s%n", currentRoutine.getRoutineName());

			for (BasicBlock currentBasicBlock : currentRoutine.getRoutine())
				for (Instruction currentInstruction : currentBasicBlock.getBlock()) {
					HashSet<Register> currentDefs = currentInstruction.getDefs();
					HashSet<Register> currentRefs = currentInstruction.getRefs();
					HashMap<Register, Register> currentDefRegReplaceMap = new HashMap<Register, Register>();
					HashMap<Register, Register> currentRefRegReplaceMap = new HashMap<Register, Register>();

					for (Register currentRefReg : currentRefs) {
						// update the rest of the instructions
						if (regMap.containsKey(currentRefReg) && (regMap.get(currentRefReg) != currentRefReg)) {
							currentRefRegReplaceMap.put(currentRefReg, regMap.get(currentRefReg));
						} else {
							// no mapping found
							regMap.put(currentRefReg, currentRefReg);
						}
					}

					for (Register currentDefReg : currentDefs) {
						if (!currentDefReg.isDirty()) {
							// first time defined
							currentDefReg.setDirty();
							regMap.put(currentDefReg, currentDefReg);
						} else {
							// split into new virtual register
							Register newReg = RegisterFactory.getNewVirtualRegister();
							regMap.put(currentDefReg, newReg);
							currentDefRegReplaceMap.put(currentDefReg, newReg);
						}
					}

					for (Register currentRefReplaceReg : currentRefRegReplaceMap.keySet())
						currentInstruction.replaceRefReg(currentRefReplaceReg, currentRefRegReplaceMap.get(currentRefReplaceReg));
					for (Register currentDefReplaceReg : currentDefRegReplaceMap.keySet())
						currentInstruction.replaceDefReg(currentDefReplaceReg, currentDefRegReplaceMap.get(currentDefReplaceReg));

					System.out.printf("%s%n", currentInstruction.toString());
				}
		}
		System.out.printf("SSA conversion completed!%n");
		return true;
	}

	private static void performLva() {
		/* live variable analysis */
		Boolean isModified = true;
		System.out.printf("%n%nLVA begins...%n");
		while (isModified) {
			System.out.printf("Starting new round...%n");
			isModified = false;	// reset isModified
			for (Routine currentRoutine : program) {
				for (BasicBlock currentBasicBlock : currentRoutine.getRoutine()) {				
					HashSet<Register> currentOldInLive = currentBasicBlock.getInLive();
					HashSet<Register> currentBlockInLive = currentBasicBlock.getOutLive();;	// recompute in-live, from out-live				
					List<Instruction> currentBlockInstructions = currentBasicBlock.getBlock();

					// compute new in-live backwards
					for (int currentInstructionIndex = currentBlockInstructions.size() - 1; currentInstructionIndex >= 0; currentInstructionIndex--) {
						// registers referenced in the block, add refs
						for (Register currentReg : currentBlockInstructions.get(currentInstructionIndex).getRefs()) {
							if (!currentBlockInLive.contains(currentReg))
								currentBlockInLive.add(currentReg);
						}
						// registers defined in the block, subtract defs
						for (Register currentReg : currentBlockInstructions.get(currentInstructionIndex).getDefs()) {
							if (currentBlockInLive.contains(currentReg))
								currentBlockInLive.remove(currentReg);
						}
					} // end of instruction traversing

					if (!currentOldInLive.equals(currentBlockInLive))
						isModified = true;

					System.out.printf("Routine: %s; BasicBlock: %s; isModified = %s%n", currentRoutine.getRoutineName(), currentBasicBlock.toString(), isModified);

					if (currentBasicBlock.setInLiveAndUpdatePrdcs(currentBlockInLive))
						isModified = true;

					System.out.printf("(setInLiveAndUpdatePrdcs?) ==> isModified = %s%n", isModified);
				}
			} 
		}
		System.out.printf("LVA Completed!%n");	/* end of LVA */

		// printing of LVA RESULTS
		System.out.printf("%n%nLVA Results printing...%n");
		for (Routine currentRoutine : program) {
			System.out.printf("%nRoutine: %s%n", currentRoutine.getRoutineName());
			for (BasicBlock currentBasicBlock : currentRoutine.getRoutine()) {
				System.out.printf("BasicBlock in-live %s%n", currentBasicBlock.getInLive());
				for (Instruction currentInstruction : currentBasicBlock.getBlock())
					System.out.printf("%s%n", currentInstruction.toString());
				System.out.printf("BasicBlock out-live %s%n", currentBasicBlock.getOutLive());

			}
		}
		System.out.printf("LVA results printing completed!%n");
	}

	private static InterferenceGraph generateInterferenceGraph() {

		InterferenceGraph rsGraph = new InterferenceGraph();

		for (String currentKey : RegisterFactory.getRegisterKeys())
			rsGraph.getNewNode(currentKey);

		System.out.printf("Building interference graph...%n");
		for (Routine currentRoutine : program) {
			for (BasicBlock currentBasicBlock : currentRoutine.getRoutine()) {
				HashSet<Register> currentInLive = currentBasicBlock.getInLive();
				for (Register currentReg1 : currentInLive)
					for (Register currentReg2 : currentInLive) {
						if (!currentReg1.equals(currentReg2)) {
							Node node1 = rsGraph.getNewNode(currentReg1.toString());
							Node node2 = rsGraph.getNewNode(currentReg2.toString());
							rsGraph.addEdge(node1, node2);
						}
					}

				HashSet<Register> currentOutLive = currentBasicBlock.getOutLive();
				for (Register currentReg1 : currentOutLive)
					for (Register currentReg2 : currentOutLive) {
						if (!currentReg1.equals(currentReg2)) {
							Node node1 = rsGraph.getNewNode(currentReg1.toString());
							Node node2 = rsGraph.getNewNode(currentReg2.toString());
							rsGraph.addEdge(node1, node2);
						}
					}
			}
		}
		System.out.printf("Interference graph build success!%n");
		rsGraph.printAdjacencies();
		return rsGraph; // return an interference graph
	}


	private static boolean checkParameters(String[] args) {
		if (args.length != 2) {
			printUsage();
			return false;
		}		
		return true;
	}

	public static void main(String[] args) {
		if (checkParameters(args) == false) System.exit(1);
		InterferenceGraph graph = buildGraph(args[0]);
		graph.colourGraph(AVAILABLE_REGISTERS);
		HashMap<String, Integer> rsColourMap = graph.yieldResult();
		System.out.printf("%n%nColoured Map: %s%n", rsColourMap.toString());
		putBackResults(rsColourMap);
	}

	private static void putBackResults(HashMap<String, Integer> rsColourMap) {
		RegisterFactory.applyColours(rsColourMap);
		// Final results printing
		System.out.printf("%n%n***FINAL Results printing...%n");
		for (Routine currentRoutine : program) {
			System.out.printf("%nRoutine: %s%n", currentRoutine.getRoutineName());
			for (BasicBlock currentBasicBlock : currentRoutine.getRoutine())
				for (Instruction currentInstruction : currentBasicBlock.getBlock())
					System.out.printf("%s%n", currentInstruction.toString());
		}
		System.out.printf("FINAL results printing completed!%n");

	}
}
